# browser_manager.py
import asyncio
import os
import uuid
from pathlib import Path
from typing import Dict, Optional
from playwright.async_api import async_playwright, BrowserContext
from logger_config import setup_logger


class BrowserConfig:
    """浏览器配置常量"""
    CACHE_BASE_DIR = Path(os.getcwd()) / "workstation_cache"
    DEFAULT_EXECUTABLE_PATH = Path(os.getcwd()) / "ms-playwright" / "chrome-win" / "chrome.exe"
    USER_AGENT = ("Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) "
                  "Chrome/58.0.3029.110 Safari/537.3")
    LAUNCH_ARGS = [
        "--disable-blink-features=AutomationControlled",
        "--disable-features=msEdgeSync,msWelcomePage"
    ]
    DEFAULT_TIMEOUT = 30000  # 30秒

    def __init__(self):
        # 判断DEFAULT_EXECUTABLE_PATH文件是否存在
        if not self.DEFAULT_EXECUTABLE_PATH.exists():
            ms_playwright_dir = Path(os.getcwd()) / "ms-playwright"
            # 搜索chromium-xxx目录
            for dir_path in ms_playwright_dir.iterdir():
                if dir_path.is_dir() and dir_path.name.startswith("chromium-"):
                    self.DEFAULT_EXECUTABLE_PATH = dir_path / "chrome-win" / "chrome.exe"
                    break


class BrowserManager:
    """浏览器实例管理器（异步）"""

    def __init__(self, browser_config: BrowserConfig = None):
        self.browser_config = browser_config or BrowserConfig()
        self._playwright = None
        self._browser_type = None
        self._contexts: Dict[str, BrowserContext] = {}
        self.logger = setup_logger("browser_manager")
        self._initialized = False

    async def initialize(self):
        """异步初始化Playwright"""
        if self._initialized:
            self.logger.warning("BrowserManager already initialized")
            return

        self._playwright = await async_playwright().start()
        self._browser_type = self._playwright.chromium
        self._initialized = True
        self.logger.info("Playwright initialized successfully")

    async def create_context(
            self,
            context_id: str = None,
            reuse_cache: bool = True,
            custom_args: Optional[list] = None
    ) -> tuple[str, BrowserContext]:
        """
        创建新的浏览器上下文
        :return: (context_id, BrowserContext)
        """
        if not context_id:
            context_id = str(uuid.uuid4())
        cache_dir = self.browser_config.CACHE_BASE_DIR / context_id

        try:
            cache_dir.mkdir(parents=True, exist_ok=True)

            context = await self._browser_type.launch_persistent_context(
                executable_path=str(self.browser_config.DEFAULT_EXECUTABLE_PATH),
                headless=False,
                user_data_dir=str(cache_dir) if reuse_cache else None,
                args=custom_args or self.browser_config.LAUNCH_ARGS,
                user_agent=self.browser_config.USER_AGENT,
                timeout=self.browser_config.DEFAULT_TIMEOUT
            )

            self._contexts[context_id] = context
            self.logger.info(f"Created new context: {context_id}")
            return context_id, context

        except Exception as e:
            self.logger.error(f"Failed to create context: {str(e)}")
            raise BrowserContextError("Context creation failed") from e

    async def get_context(self, context_id: str) -> BrowserContext:
        """获取已存在的浏览器上下文"""
        if context := self._contexts.get(context_id):
            return context
        raise BrowserContextError(f"Context not found: {context_id}")

    async def close_context(self, context_id: str, remove_cache: bool = False):
        """安全关闭浏览器上下文"""
        try:
            if context := self._contexts.pop(context_id, None):
                try:
                    # 尝试访问页面列表触发异常（用于检测上下文状态）
                    _ = context.pages
                    await context.close()  # 直接尝试关闭
                    self.logger.info(f"Closed context: {context_id}")
                except Exception as e:
                    if "been closed" in str(e):
                        self.logger.debug(f"上下文已关闭: {context_id}")  # 改为调试级别日志
                    else:
                        self.logger.error(f"关闭上下文异常: {str(e)}")
                        raise

                self.logger.info(f"Closed context: {context_id}")

                if remove_cache:
                    cache_dir = self.browser_config.CACHE_BASE_DIR / context_id
                    if cache_dir.exists():
                        import shutil
                        shutil.rmtree(cache_dir)
                        self.logger.info(f"Removed cache for: {context_id}")
        except Exception as e:
            self.logger.error(f"Error closing context {context_id}: {str(e)}")
            raise BrowserContextError(f"Failed to close context: {context_id}") from e

    async def close_all_contexts(self):
        """关闭所有浏览器上下文"""
        closing_tasks = [
            self.close_context(ctx_id)
            for ctx_id in list(self._contexts.keys())
        ]
        await asyncio.gather(*closing_tasks)
        self.logger.info("All browser contexts closed")

    async def shutdown(self):
        """完全关闭Playwright"""
        if self._initialized:
            await self.close_all_contexts()
            await self._playwright.stop()
            self._initialized = False
            self.logger.info("Playwright shutdown complete")


class BrowserContextError(Exception):
    """浏览器上下文操作异常"""

    def __init__(self, message: str, context_id: Optional[str] = None):
        super().__init__(message)
        self.context_id = context_id
        self.message = f"[{context_id}] {message}" if context_id else message

    def __str__(self):
        return self.message
